/*
 * Copyright (C) 2013, Pavel Samko <bulldozerbsg@gmail.com>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of the "Pavel Samko" nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
 * OF SUCH DAMAGE.
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
#include <stdint.h>

#if defined _WIN32
#include <windows.h>
#else
#include <sys/types.h>
#include <dirent.h>
#include <unistd.h>
#endif

#include "vfsformat.h"
#include "vfscrc16.h"
#include "vfsutils.h"
#include "vfscrypttea.h"
#include "vfspacklzma.h"

#define VER "v0.9"

#define VFS_PATH_MAX 256
#define VFS_VERSION 1

char * g_out_file_path = NULL;
char * g_in_dir_path = NULL;
int g_error_count = 0;
int g_verbose = 0;
int g_tea = 0;
uint32_t g_tea_key[4] = {0};
char * g_tea_file_path = NULL; // C header file for TEA key
int g_lzma = 0;

VIRTFS_HEADER g_header = {{VIRTFS_MAGIC_DATA}, VFS_VERSION, 0, 0};

typedef struct VFSDIR
{
#if defined _WIN32
	HANDLE m_find;
	WIN32_FIND_DATA m_find_data;
	char m_has_data;
#else
	DIR * m_dir;
	struct dirent * m_dirent;
#endif
} VFSDIR;

#define DIR_TYPE_FILE 0
#define DIR_TYPE_DIR 1

char * vfsgetwd(char * path, size_t path_size)
{
	if (path)
	{
#if defined _WIN32
		GetCurrentDirectoryA(path_size, path);
#else
		getcwd(path, path_size);
#endif
	}
	return path;
}

void vfssetwd(const char * path)
{
	if (path)
	{
#if defined _WIN32
		SetCurrentDirectoryA(path);
#else
		chdir(path);
#endif
	}
}

typedef struct VFSDIRENT
{
	char m_type;
	char m_name[VFS_PATH_MAX];
} VFSDIRENT;

VFSDIR * dopen(char * name)
{
	VFSDIR * result = (VFSDIR *)malloc(sizeof(VFSDIR));
	if (result && name)
	{
		memset(result, 0, sizeof(VFSDIR));
#if defined _WIN32
		vfssetwd(name);
		result->m_has_data = 0;
		result->m_find = FindFirstFileA("*", &result->m_find_data);
		if (result->m_find != INVALID_HANDLE_VALUE)
		{
			result->m_has_data = 1;
		}
#else
		result->m_dir = opendir(name);
#endif
	}
	return result;
}

void dclose(VFSDIR * dir)
{
	if (dir)
	{
#if defined _WIN32
		if (dir->m_find != INVALID_HANDLE_VALUE)
		{
			FindClose(dir->m_find);
		}
#else
		if (dir->m_dir)
		{
			closedir(dir->m_dir);
			dir->m_dir = NULL;
		}
#endif
		free(dir);
	}
}

int dread(VFSDIRENT * ent, VFSDIR * dir)
{
	int result = 0;
#if defined _WIN32
	if (dir && ent && dir->m_has_data && dir->m_find != INVALID_HANDLE_VALUE)
	{
		memset(ent, 0, sizeof(VFSDIRENT));
		ent->m_type = dir->m_find_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY ? DIR_TYPE_DIR : DIR_TYPE_FILE;
		strcpy(ent->m_name, dir->m_find_data.cFileName);
		dir->m_has_data = FindNextFileA(dir->m_find, &dir->m_find_data) == TRUE ? 1 : 0;
		result = 1;
	}
#else
	if (dir && ent)
	{
dread_loop:
		dir->m_dirent = readdir(dir->m_dir);
		if (dir->m_dirent)
		{
			memset(ent, 0, sizeof(VFSDIRENT));
			if (dir->m_dirent->d_type == DT_DIR)
			{
				ent->m_type = DIR_TYPE_DIR;
			} else
			if (dir->m_dirent->d_type == DT_REG)
			{
				ent->m_type = DIR_TYPE_FILE;
			} else
			{
				goto dread_loop; // unsuported type skeep
			}
			strcpy(ent->m_name, dir->m_dirent->d_name);
			result = 1;
		}
	}
#endif
	return result;
}

#if defined _WIN32
#define DIR_DELIMITER_S "\\"
#else
#define DIR_DELIMITER_S "/"
#endif

typedef struct LISTREC LISTREC;
struct LISTREC
{
	VIRTFS_RECORD m_rec;
	VIRTFS_NAMEHASH m_namehash;
	VIRTFS_PACKTABLE m_packtable;
	VIRTFS_PACKTABLE_RECORD * m_packtable_recs;
	LISTREC * m_next;
};

LISTREC * create_list(char * base_path)
{
	LISTREC * head = NULL;
	LISTREC * tmp;
	LISTREC * tmp2;
	char * curdir = (char *)malloc(VFS_PATH_MAX);
	char * path;
	VFSDIR * dir;
	VFSDIRENT ent;

	memset(curdir, 0, VFS_PATH_MAX);
	vfsgetwd(curdir, VFS_PATH_MAX);
	dir = dopen(curdir);
	while (dread(&ent, dir))
	{
		if (ent.m_type == DIR_TYPE_DIR)
		{
			if (!strcmp(ent.m_name, ".") || !strcmp(ent.m_name, ".."))
			{
				continue;
			}
			vfssetwd(ent.m_name);
			tmp = create_list(base_path);
			if (tmp)
			{
				tmp2 = tmp;
				// find last
				while (tmp2->m_next)
				{
					tmp2 = tmp2->m_next;
				}
				tmp2->m_next = head;
				head = tmp;
			}
			vfssetwd(curdir);
		} else
		if (ent.m_type == DIR_TYPE_FILE)
		{
			tmp = (LISTREC *)malloc(sizeof(LISTREC));
			memset(tmp, 0, sizeof(LISTREC));
			path = curdir + strlen(base_path);
			if (strlen(path))
			{
				sprintf(tmp->m_rec.m_name, "%s%s%s", path + 1, DIR_DELIMITER_S, ent.m_name);
			} else
			{
				sprintf(tmp->m_rec.m_name, "%s", ent.m_name);
			}
			tmp->m_rec.m_name_size = strlen(tmp->m_rec.m_name);
			if (g_tea)
			{
				tmp->m_rec.m_crypt_format = VIRTFS_CRYPT_TEA;
			}
			if (g_lzma)
			{
				tmp->m_rec.m_compress_format = VIRTFS_COMPRESS_LZMA;
			}
			tmp->m_next = head;
			head = tmp;
		}
	}
	dclose(dir);
	free(curdir);
	return head;
}

void clean_list(LISTREC * list)
{
	LISTREC * tmp;
	while (list)
	{
		if (list->m_packtable_recs)
		{
			free(list->m_packtable_recs);
		}
		tmp = list->m_next;
		free(list);
		list = tmp;
	}
}

void write_list(LISTREC * list, FILE * f)
{
	while (list)
	{
		list->m_namehash.m_offset = ftell(f);
		list->m_namehash.m_crc16_ccitt = vfs_crc16_ccitt(VFS_CRC16_CCITT_INIT, (uint8_t *)list->m_rec.m_name, list->m_rec.m_name_size);
		fwrite(&list->m_rec, sizeof(VIRTFS_RECORD) - VIRTFS_NAME_SIZE, 1, f);
		fwrite(list->m_rec.m_name, list->m_rec.m_name_size, 1, f);
		list = list->m_next;
	}
}

void write_list_data(LISTREC * list, FILE * f)
{
	FILE * fdata;
	size_t buff_pack_size = VFS_LZMA_BLOCKSIZE * 1.3 + 128;
	char * buff = (char *)malloc(VFS_LZMA_BLOCKSIZE);
	char * buff_pack = (char *)malloc(buff_pack_size);
	size_t read_size, packed_size_piece;
	size_t need_read_size, need_crypt_size;
	size_t i, j, t, c;
	size_t packed_size;
	while (list)
	{
		fdata = fopen(list->m_rec.m_name, "rb");
		if (fdata)
		{
			if (g_verbose)
			{
				printf("Pack %s\n", list->m_rec.m_name);
			}
			list->m_rec.m_offset = ftell(f);
			// get file size
			fseek(fdata, 0, SEEK_END);
			need_read_size = ftell(fdata);
			fseek(fdata, 0, SEEK_SET);
			if (g_lzma)
			{
				packed_size = 0;
				list->m_packtable.m_rec_count = ((need_read_size + VFS_LZMA_BLOCKSIZE - 1) / VFS_LZMA_BLOCKSIZE);
				j = list->m_packtable.m_rec_count * sizeof(VIRTFS_PACKTABLE_RECORD);
				list->m_packtable_recs = (VIRTFS_PACKTABLE_RECORD *)malloc(j);
				memset(list->m_packtable_recs, 0, j);
				fwrite(&list->m_packtable, sizeof(VIRTFS_PACKTABLE), 1, f);
				fwrite(list->m_packtable_recs, j, 1, f);
				j = 0;
				c = ftell(f);
				do
				{
					packed_size_piece = 0;
					read_size = need_read_size < VFS_LZMA_BLOCKSIZE ? need_read_size : VFS_LZMA_BLOCKSIZE;
					read_size = fread(buff, 1, read_size, fdata);
					if (read_size)
					{
						list->m_packtable_recs[j].m_offset = ftell(f); // offset
						packed_size_piece = vfs_lzma_pack(buff_pack, buff_pack_size, buff, read_size);
						list->m_packtable_recs[j++].m_size = packed_size_piece; // size
					}
					fwrite(buff_pack, 1, packed_size_piece, f);
					need_read_size -= read_size;
					list->m_rec.m_size += read_size;
					packed_size += packed_size_piece;
				} while (read_size);
				if (g_tea)
				{
					t = ftell(f);
					need_crypt_size = (t - c);
					fseek(f, c, SEEK_SET);
					do
					{
						read_size = need_crypt_size < VFS_TEA_BLOCKSIZE ? need_crypt_size : VFS_TEA_BLOCKSIZE;
						read_size = fread(buff, 1, read_size, f);
						for (i = read_size; i < VFS_TEA_BLOCKSIZE; ++i)
						{
							buff[i] = 0x01;
						}
						vfs_tea_encrypt(buff, VFS_TEA_BLOCKSIZE, g_tea_key);
						if (read_size)
						{
							fseek(f, -read_size, SEEK_CUR);
							fwrite(buff, 1, VFS_TEA_BLOCKSIZE, f);
							fseek(f, 0, SEEK_CUR); /// FIXED BUG: not read data after first write (previous line)
						}
						need_crypt_size -= read_size;
					} while (read_size);
				}
				t = ftell(f);
				fseek(f, list->m_rec.m_offset, SEEK_SET);
                j = list->m_packtable.m_rec_count * sizeof(VIRTFS_PACKTABLE_RECORD);
				fwrite(&list->m_packtable, sizeof(VIRTFS_PACKTABLE), 1, f);
				fwrite(list->m_packtable_recs, j, 1, f);
				fseek(f, t, SEEK_SET);
			} else
			{
				if (g_tea)
				{
					do
					{
						read_size = need_read_size < VFS_TEA_BLOCKSIZE ? need_read_size : VFS_TEA_BLOCKSIZE;
						read_size = fread(buff, 1, read_size, fdata);
						for (i = read_size; i < VFS_TEA_BLOCKSIZE; ++i)
						{
							buff[i] = 0x01;
						}
						vfs_tea_encrypt(buff, VFS_TEA_BLOCKSIZE, g_tea_key);
						if (read_size)
						{
							fwrite(buff, 1, VFS_TEA_BLOCKSIZE, f);
						}
						list->m_rec.m_size += read_size;
						need_read_size -= read_size;
					} while (read_size);
				} else
				{
					do
					{
						read_size = need_read_size < VFS_LZMA_BLOCKSIZE ? need_read_size : VFS_LZMA_BLOCKSIZE;
						read_size = fread(buff, 1, read_size, fdata);
						fwrite(buff, 1, read_size, f);
						list->m_rec.m_size += read_size;
						need_read_size -= read_size;
					} while (read_size);
				}
			}
			fclose(fdata);
		} else
		{
			printf("[ERROR] Could not open a file %s\n", list->m_rec.m_name);
			++g_error_count;
		}
		list = list->m_next;
	}
	free(buff);
	free(buff_pack);
}

void write_list_hash(LISTREC * list, FILE * f)
{
	while (list)
	{
		fwrite(&list->m_namehash, sizeof(VIRTFS_NAMEHASH), 1, f);
		list = list->m_next;
	}
}

unsigned short get_list_count(LISTREC * list)
{
	unsigned short result = 0;
	while (list)
	{
		++result;
		list = list->m_next;
	}
	return result;
}

void list_change_slash(LISTREC * list)
{
	while (list)
	{
		vfs_strbstos(list->m_rec.m_name);
		list = list->m_next;
	}
}

void parse_tea_key(const char * key)
{
	char * keydup = vfs_strdup(key);
	char * pch = strtok(keydup, " \",");
	int count = 0;
	uint32_t value = 0;
	while (pch != NULL)
	{
		sscanf(pch, "0x%x", &value);
		if (count < 4)
		{
			g_tea_key[count] = value;
		}
		count++;
		pch = strtok(NULL, " \",");
	}
	free(keydup);
}

void gen_tea_key(void)
{
	srand((int)time(NULL));
	if (!g_tea_key[0] && !g_tea_key[1] && !g_tea_key[2] && !g_tea_key[3])
	{
		g_tea_key[0] = rand();
		g_tea_key[1] = rand();
		g_tea_key[2] = rand();
		g_tea_key[3] = rand();
	}
}

int main(int argc, char * argv[])
{
	FILE * fout;
	LISTREC * list;
	char * base_path;
	int i;

	printf("VFS Packer " VER " - BulldozerBSG (c) 2013\n\n");
	// parse cmd
	for (i = 1; i < argc; ++i)
	{
		if (!strcmp(argv[i], "-o"))
		{
			++i;
			if (i < argc)
			{
				g_out_file_path = argv[i];
			}
		} else
		if (!strcmp(argv[i], "-c"))
		{
			g_tea = 1;
		} else
		if (!strcmp(argv[i], "-k"))
		{
			++i;
			if (i < argc)
			{
				parse_tea_key(argv[i]);
			}
		} else
		if (!strcmp(argv[i], "-h"))
		{
			++i;
			if (i < argc)
			{
				g_tea_file_path = argv[i];
			}
		} else
		if (!strcmp(argv[i], "-p"))
		{
			g_lzma = 1;
		} else
		if (!strcmp(argv[i], "--verbose") || !strcmp(argv[i], "-v"))
		{
			g_verbose = 1;
		} else
		if (!strcmp(argv[i], "--help") || !strcmp(argv[i], "-?"))
		{
			printf("Usage: vfspack [params] [directory]\n");
			printf("Params:\n");
			printf("  -o [file]       Output file name with archive? default \"out.vfs\"\n");
			printf("  -c              Encrypt using TEA\n");
			printf("  -k [128b key]   TEA key (example \"0x234590ba, 0xfc279847, 0x56237882, 0x12238a91\"\n");
			printf("  -h [file]       Generate C header file with TEA key\n");
			printf("  -p              Pack using LZMA\n");
			printf("  -v\n");
			printf("  --verbose       Print all info, default only errors\n");
			printf("  -?\n");
			printf("  --help          Print this help\n\n");
			return EXIT_SUCCESS;
		} else
		{
			g_in_dir_path = argv[i];
		}
	}
	// TEA gen key
	if (g_tea)
	{
		gen_tea_key();
		if (g_tea_file_path)
		{
			fout = fopen(g_tea_file_path, "wb");
			if (fout)
			{
				fprintf(fout, "#ifndef VFS_TEA_KEY_H\n");
				fprintf(fout, "#define VFS_TEA_KEY_H\n");
				fprintf(fout, "\n");
				fprintf(fout, "#define VFS_TEA_KEY {0x%x, 0x%x, 0x%x, 0x%x}\n", g_tea_key[0], g_tea_key[1], g_tea_key[2], g_tea_key[3]);
				fprintf(fout, "\n");
				fprintf(fout, "#endif\n");
				fflush(fout);
				fclose(fout);
			}
		}
	}
	// pack
	fout = fopen(g_out_file_path ? g_out_file_path : "out.vfs", "wb+");
	if (fout)
	{
		if (g_in_dir_path)
		{
			vfssetwd(g_in_dir_path);
		}
		// parse dir & create list
		base_path = (char *)malloc(VFS_PATH_MAX);
		memset(base_path, 0, VFS_PATH_MAX);
		vfsgetwd(base_path, VFS_PATH_MAX);
		list = create_list(base_path);
		free(base_path);
		// vrite header
		fwrite(&g_header, sizeof(VIRTFS_HEADER), 1, fout);
		// write records from list
		write_list(list, fout);
		// write data from list
		write_list_data(list, fout);
		// write hash
		g_header.m_name_hash_table_offset = ftell(fout);
		write_list_hash(list, fout);
		// get count
		g_header.m_rec_count = get_list_count(list);
		fseek(fout, 0, SEEK_SET);
		// vrite header with count
		fwrite(&g_header, sizeof(VIRTFS_HEADER), 1, fout);
		// write records from list with offsets and size and slash
		list_change_slash(list);
		write_list(list, fout);
		fseek(fout, 0, SEEK_END);
		fclose(fout);
		clean_list(list);

		if (g_error_count)
		{
			printf("\nPacked end with errors, errors count = %d\n\n", g_error_count);
			return EXIT_FAILURE;
		} else
		{
			printf("\nPacked end without errors\n\n");
		}
	} else
	{
	}
	return EXIT_SUCCESS;
}
